
import matplotlib.pyplot as plt

import numpy as np
import pandas as pd

import os, sys


from tqdm import tqdm
import json
import numpy as np
from sklearn.preprocessing import LabelEncoder
import category_encoders

import csv
from sklearn.metrics import roc_auc_score

datalib_path = os.path.abspath('..')
sys.path.append(datalib_path)

train = pd.read_csv(datalib_path+'/raw_data/train_dataset.csv', delimiter="\t")
test = pd.read_csv(datalib_path+'/raw_data/test_dataset.csv', delimiter="\t")



data = pd.concat([train, test])


data['location_first_lvl'] = data['location'].astype(str).apply(lambda x: json.loads(x)['first_lvl'])
data['location_sec_lvl'] = data['location'].astype(str).apply(lambda x: json.loads(x)['sec_lvl'])
data['location_third_lvl'] = data['location'].astype(str).apply(lambda x: json.loads(x)['third_lvl'])


data['auth_type'].fillna('__NaN__', inplace=True)

data['op_date'] = pd.to_datetime(data['op_date'])
data['op_ts'] = data["op_date"].values.astype(np.int64) // 10 ** 11
data.drop(['session_id','op_date'], axis=1, inplace=True)


for col in tqdm(['user_name', 'action', 'auth_type', 'ip', 'client_type', 'browser_source',
                 'ip_location_type_keyword', 'ip_risk_level', 'location', 'device_model',
                 'os_type', 'os_version', 'browser_type', 'browser_version',
                 'bus_system_code', 'op_target', 'location_first_lvl', 'location_sec_lvl', 
                 'location_third_lvl']):

    ce=category_encoders.count.CountEncoder()
    data[col] = ce.fit_transform(data[col])

    
data = data.sort_values(by=['user_name', 'op_ts']).reset_index(drop=True)
data['last_ts'] = data.groupby(['user_name'])['op_ts'].shift(1)
data['last_ts'].fillna(value=0, inplace=True)
data['ts_diff1'] = data['op_ts'] - data['last_ts']

for f in ['ip', 'location', 'device_model', 'os_version', 'browser_version']:
    data[f'user_{f}_nunique'] = data.groupby(['user_name'])[f].transform('nunique')

for method in ['mean', 'max', 'min', 'std','sum','median']:
    for col in ['user_name','ip', 'location', 'device_model', 'os_version', 'browser_version']:
        data[f'ts_diff1_{method}_'+str(col)] = data.groupby(col)['ts_diff1'].transform(method)
        

train = data[data['risk_label'].notna()]
test = data[data['risk_label'].isna()]
y=train["risk_label"].copy()
train.drop("risk_label" ,axis= 1 ,inplace=True)
test.drop("risk_label" ,axis= 1 ,inplace=True)

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(train, y, random_state=1)


import lightgbm as lgbm
from sklearn.metrics import roc_auc_score

train_data = lgbm.Dataset(X_train, y_train)
val_data = lgbm.Dataset(X_test, y_test)

# define parameters
parameters = {
    'objective': 'binary',
    'metric': 'auc',
    'is_unbalance': 'true',
    'max_depth': 4,
    'learning_rate':0.1,
    'feature_fraction': 0.5,
    'bagging_fraction': 0.4,
    'bagging_freq': 5,
    'num_threads' : 4,
    'seed' : 76
}

# train lightGBM model
model = lgbm.train(parameters,
                   train_data,
                   num_boost_round=150,
                   valid_sets=[train_data, val_data]
                  )
            

import catboost
cat_columns = ['user_name', 'action', 'auth_type', 'ip', 'client_type', 'browser_source',
                 'ip_location_type_keyword', 'ip_risk_level', 'location', 'device_model',
                 'os_type', 'os_version', 'browser_type', 'browser_version',
                 'bus_system_code', 'op_target', 'location_first_lvl', 'location_sec_lvl', 
                 'location_third_lvl']

cb = catboost.CatBoostClassifier(cat_features = cat_columns, eval_metric='AUC', random_state = 1988)
cb.fit(X_train, y_train, verbose=200, eval_set=(X_test,y_test))

pred_lgb=model.predict(X_test)
pred_cb=cb.predict_proba(X_test)[:,1]
print(roc_auc_score(y_test,pred_lgb))
print(roc_auc_score(y_test,pred_cb))

pred_lgb=1-model.predict(test) #lgb overfitting leads to worse predictions
pred_cb=cb.predict_proba(test)[:,1]


predictions = zip(list(range(1,len(pred_lgb)+1)), pred_lgb)
with open(datalib_path+"/prediction_result/submission_lgb.csv","w",newline ='') as pred:
    csv_out = csv.writer(pred)
    csv_out.writerow(['id','ret'])
    for row in predictions:
        csv_out.writerow(row)

predictions = zip(list(range(1,len(pred_lgb)+1)), pred_cb)
with open(datalib_path+"/prediction_result/submission_cb.csv","w",newline ='') as pred:
    csv_out = csv.writer(pred)
    csv_out.writerow(['id','ret'])
    for row in predictions:
        csv_out.writerow(row)
